#!/usr/bin/python
# -*- coding:utf-8 -*-
import base64
import io
import hashlib
import hmac
import struct
import time
from os import urandom
from binascii import unhexlify, hexlify
from urllib.parse import quote, urlencode
import qrcode

TOTP_LEN = 6
COMPANY_NAME = 'knownsec.com'


def hotp(key, counter, length):
    """
    hmac otp generator
    """
    msg = struct.pack(">Q", counter)
    hs = hmac.new(key, msg, hashlib.sha1).digest()
    offset = hs[19] & 0x0F
    bin_code = struct.unpack(">I", hs[offset : offset + 4])[0] & 0x7FFFFFFF
    hotp = bin_code % pow(10, length)
    hotp = str(hotp)
    while len(hotp) < length:
        hotp = '0' + hotp
    return hotp


class TOTP(object):
    """
    totp interface
    """

    def __init__(self, key, step=30, t0=0, length=6, drift=0):
        self.key = key
        self.step = step
        self.t0 = t0
        self.length = length
        self.drift = drift
        self._time = None

    def t(self):
        return ((int(self.time) - self.t0) // self.step) + self.drift

    def token(self):
        return hotp(self.key, self.t(), length=self.length)

    @property
    def time(self):
        return self._time if self._time else time.time()

    @time.setter
    def time(self, value):
        self._time = value

    @time.deleter
    def time(self):
        self._time = None

    def verify(self, token, tolerance=0, min_t=None):
        """
        token: the provided token
        tolerance: the amount of clock drift you're willing to accommdate in steps
        min_t: the mininum t value we'll accept
        """
        drift_orig = self.drift
        verified = False
        for offset in range(-tolerance, tolerance + 1):
            self.drift = drift_orig + offset
            if min_t and self.t() < min_t:
                continue
            elif self.token() == token:
                verified = True
                break
        else:
            self.drift = drift_orig
        return verified


def totp(key, step=30, t0=0, length=6, drift=0):
    return TOTP(key, step, t0, length, drift).token()


def random_hex(length=20):
    return hexlify(urandom(length))


def unhexlify_key(key):
    return unhexlify(key.encode())


def totp_len():
    return TOTP_LEN or 6


def get_otpauth_url(username, secret, issuer=None, length=None):
    username = username.encode('utf8')
    issuer = issuer.encode('utf8') if issuer else None
    label = quote(b': '.join([issuer, username]) if issuer else username)
    query = [('secret', secret), ('digits', length or totp_len())]
    if issuer:
        query.append(('issuer', issuer))

    return 'otpauth://totp/%s?%s' % (label, urlencode(query))


def create_totp_qr_secret(key=None):
    """
    create totp verify secret and qr secret
    """
    if not key:
        key = random_hex(20).decode('ascii')
    rawkey = unhexlify(key.encode('ascii'))
    b32key = base64.b32encode(rawkey).decode('utf-8')
    return key, b32key


def create_qr_url(username, secret):
    """
    create a image for google authenticator scan
    """
    otpauth_url = get_otpauth_url(username, secret, COMPANY_NAME, length=totp_len())
    qr = qrcode.QRCode(
        version=5,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=8,
        border=4,
    )
    qr.add_data(otpauth_url)
    qr.make(fit=True)
    img = qr.make_image()
    img_buffer = io.BytesIO()
    img.save(img_buffer, 'png')
    res = img_buffer.getvalue()
    img_buffer.close()
    return base64.b64encode(res)


def verify_token(key, token, tolerance=2):
    """
    verify token, tolerance is 2 step, yet is one minute
    """
    totp = TOTP(key)
    validated = totp.verify(token, tolerance)
    return validated


if __name__ == '__main__':
    pass
